/*******************************************************************************
* file  : rsa key text
* creat : apleilx
* data  : 2017/03/02
*******************************************************************************/
#include "common.h"
#include "sys_task.h"
#include "os_obj.h"
#include "mbedtls/rsa.h"
#include "mbedtls/entropy.h"
#include "mbedtls/ctr_drbg.h"
#include "mbedtls/platform.h"
#include "mbedtls/bignum.h"

extern uu8 * rsa_key_N;
extern uu8 * rsa_key_E;
extern uu8 * rsa_key_D;
extern uu8 * rsa_key_P;
extern uu8 * rsa_key_Q;

extern void * rsa_src, *rsa_enc, *rsa_dec;

void ras_ext_enc(void);
void ras_ext_dec(void);

#define DEV_RANDOM_THRESHOLD        32
int rsa_random_entropy_poll( void *, unsigned char *, size_t, size_t *);
static const char pers[] = "simple_rsa";

/*
********************************************************************************
*                              assert
*******************************************************************************/
#define assert_exit(cond, ret) \
    do { if (!(cond)) { \
        dbg_print("  !. assert: failed [line: %d, error: -0x%04X]\n", __LINE__, -ret); \
        goto cleanup; \
    } } while (0)

/*******************************************************************************
* @brief  pub dec.
* \param[in] none
* \retval: none
*******************************************************************************/
void ras_pub_enc(void)
{
	dbg_print("\n\n");
	dbg_print("======================================\n");
	dbg_print("public key encode, privite key decode \n");
	dbg_print("======================================\n");
	
	ras_ext_enc();
	ras_ext_dec();
}


void ras_ext_enc(void)
{
	int ret;
	
    mbedtls_rsa_context *ctx;
    mbedtls_entropy_context *entropy;
    mbedtls_ctr_drbg_context *ctr_drbg;
	
	
	entropy = malloc(sizeof(mbedtls_entropy_context));
	ctr_drbg = malloc(sizeof(mbedtls_ctr_drbg_context));
	ctx = malloc(sizeof(mbedtls_rsa_context));

	memset(rsa_src, 0, 100);
	loop_n(20) ((uu16 *)(rsa_src))[index] = index;
	
    //rsa结构体初始化
    mbedtls_rsa_init(ctx);  
    mbedtls_entropy_init(entropy);
    mbedtls_ctr_drbg_init(ctr_drbg);
	
	// 填充方案OAEP //SHA256做散列算法
    mbedtls_rsa_set_padding(ctx, MBEDTLS_RSA_PKCS_V21, MBEDTLS_MD_SHA256);
	
	// 随机数生成器
	ret = mbedtls_entropy_add_source( 
		entropy, 
		rsa_random_entropy_poll,
		NULL, 
		DEV_RANDOM_THRESHOLD,
		MBEDTLS_ENTROPY_SOURCE_STRONG );
	
	assert_exit(ret == 0, ret);
	
    ret = mbedtls_ctr_drbg_seed(
		ctr_drbg, 
		mbedtls_entropy_func, 
		entropy, 
        (const uint8_t *) pers, 
		strlen(pers));				//根据个性化字符串更新种子
	
    assert_exit(ret == 0, ret);
	
	// 导入KEY
	ret = mbedtls_rsa_import_raw(ctx, 
		rsa_key_N, 1024/8,
		NULL, 1024/8/2,
		NULL, 1024/8/2,
		NULL, 1024/8,
		rsa_key_E, 16
		);
		
	ret += mbedtls_rsa_complete(ctx);
	assert_exit(ret == 0, ret);
		
    // RSA 公匙 加密
	// 消息长度应小于 （密匙长度-11）
	// 密文长度和密匙长度相同
    ret = mbedtls_rsa_pkcs1_encrypt(
		ctx, 
		mbedtls_ctr_drbg_random,	//随机数生成接口
		ctr_drbg,          			//随机数结构体
		40,							//消息长度
		rsa_src,					//输入消息指针
		rsa_enc);					//输出密文指针
    assert_exit(ret == 0, ret);                              

cleanup:
    mbedtls_ctr_drbg_free(ctr_drbg);	//释放随机数结构体
    mbedtls_entropy_free(entropy);		//释放熵结构体
    mbedtls_rsa_free(ctx);				//释放rsa结构体

	// 释放内存
	free(ctx);
	free(entropy);
	free(ctr_drbg);
}

void ras_ext_dec(void)
{
	int ret;
	uu32 olen;
	
    mbedtls_rsa_context *ctx;
    mbedtls_entropy_context *entropy;
    mbedtls_ctr_drbg_context *ctr_drbg;
	
	
	entropy = malloc(sizeof(mbedtls_entropy_context));
	ctr_drbg = malloc(sizeof(mbedtls_ctr_drbg_context));
	ctx = malloc(sizeof(mbedtls_rsa_context));
	
    //rsa结构体初始化
    mbedtls_rsa_init(ctx);  
    mbedtls_entropy_init(entropy);
    mbedtls_ctr_drbg_init(ctr_drbg);
	
	// 填充方案OAEP //SHA256做散列算法
    mbedtls_rsa_set_padding(ctx, MBEDTLS_RSA_PKCS_V21, MBEDTLS_MD_SHA256);
	
	// 随机数生成器
	ret = mbedtls_entropy_add_source( 
		entropy, 
		rsa_random_entropy_poll,
		NULL, 
		DEV_RANDOM_THRESHOLD,
		MBEDTLS_ENTROPY_SOURCE_STRONG );
	
	assert_exit(ret == 0, ret);
	
    ret = mbedtls_ctr_drbg_seed(
		ctr_drbg, 
		mbedtls_entropy_func, 
		entropy, 
        (const uint8_t *) pers, 
		strlen(pers));				//根据个性化字符串更新种子
	
    assert_exit(ret == 0, ret);
	
	// 导入KEY
	ret = mbedtls_rsa_import_raw(ctx, 
		rsa_key_N, 1024/8,
		rsa_key_P, 1024/8/2,
		rsa_key_Q, 1024/8/2,
		rsa_key_D, 1024/8,
		rsa_key_E, 16
		);
		
	ret += mbedtls_rsa_complete(ctx);
	assert_exit(ret == 0, ret);
		
	// RSA 私匙 解密
    ret = mbedtls_rsa_pkcs1_decrypt(
		ctx, 
		mbedtls_ctr_drbg_random,	//随机数生成接口
		ctr_drbg,					//随机数结构体
		&olen,						//输出长度
		rsa_enc,					//输入密文指针
		rsa_dec,					//输出明文指针
		2048);						//最大输出明文数组长度
    assert_exit(ret == 0, ret);                                 
	
	dbg_u16_print("RSA ENC:\n",rsa_enc,1024/16,16);
	dbg_u16_print("RSA DEC:\n",rsa_dec,olen/2,16);
	
cleanup:
    mbedtls_ctr_drbg_free(ctr_drbg);	//释放随机数结构体
    mbedtls_entropy_free(entropy);		//释放熵结构体
    mbedtls_rsa_free(ctx);				//释放rsa结构体

	// 释放内存
	free(ctx);
	free(entropy);
	free(ctr_drbg);
}
